from parl.env.atari_wrappers import FrameStack

OBS_FORMAT = 'NCHW'


class Env:
    def __init__(self, env, env_seed=None, test=False, ob_rms=None):
        self.env = FrameStack(env, 4, OBS_FORMAT)

        self.obs_space = self.env.observation_space
        self.act_space = self.env.action_space

    def reset(self):
        return self.env.reset()

    def step(self, action):
        return self.env.step(action)

    def render(self):
        return self.env.render()
